#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from torch_geometric.datasets import ModelNet
from torch_geometric.transforms import SamplePoints, NormalizeScale
from torch_geometric.loader import DataLoader

def load_modelnet_dataset():
    """
    加载ModelNet数据集（使用PyTorch Geometric）
    
    Returns:
        训练数据集和测试数据集
    """
    # 定义转换：采样点并归一化
    # SamplePoints：将网格转换为点云，随机采样固定数量的点
    # NormalizeScale：将点云归一化到单位球内
    transform = [SamplePoints(1024), NormalizeScale()]
    
    # 加载ModelNet40数据集
    train_dataset = ModelNet(
        root="./data", 
        name='40',
        train=True, 
        transform=transform
    )
    
    test_dataset = ModelNet(
        root="./data", 
        name='40',
        train=False, 
        transform=transform
    )
    
    print(f"ModelNet40 训练集大小: {len(train_dataset)}")
    print(f"ModelNet40 测试集大小: {len(test_dataset)}")
    print(f"类别数量: {train_dataset.num_classes}")
    
    return train_dataset, test_dataset

def visualize_modelnet_samples(dataset, num_samples=3):
    """
    可视化ModelNet数据集中的样本
    
    Args:
        dataset: ModelNet数据集
        num_samples: 要可视化的样本数量
    """
    fig = plt.figure(figsize=(15, 5 * num_samples))
    
    for i in range(num_samples):
        # 随机选择一个样本
        idx = np.random.randint(0, len(dataset))
        data = dataset[idx]
        
        # 获取点云坐标和类别
        points = data.pos.numpy()  # 点云坐标
        category = dataset.y_mask[data.y.item()]  # 类别名称
        
        # 创建3D图
        ax = fig.add_subplot(num_samples, 1, i+1, projection='3d')
        
        # 绘制点云
        ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=2, c=points[:, 2], cmap='viridis')
        
        # 设置坐标轴标签
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        
        # 设置标题
        ax.set_title(f'类别: {category}')
        
        # 设置视角
        ax.view_init(elev=30, azim=45)
    
    plt.tight_layout()
    plt.show()

def create_dataloader(dataset, batch_size=32):
    """
    创建数据加载器
    
    Args:
        dataset: ModelNet数据集
        batch_size: 批次大小
        
    Returns:
        数据加载器
    """
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=4
    )

def print_batch_info(batch):
    """
    打印批次信息
    
    Args:
        batch: 数据批次
    """
    print(f"批次形状:")
    print(f"- 点云坐标: {batch.pos.shape}")
    print(f"- 点云特征 (如果有): {batch.x.shape if hasattr(batch, 'x') else 'None'}")
    print(f"- 类别标签: {batch.y.shape}")
    print(f"- 批次大小: {batch.num_graphs}")
    print(f"- 样本中的点数: {batch.num_nodes}")
    
    # 打印点云坐标的统计信息
    print(f"\n点云坐标统计:")
    print(f"- 最小值: {batch.pos.min(dim=0)[0]}")
    print(f"- 最大值: {batch.pos.max(dim=0)[0]}")
    print(f"- 均值: {batch.pos.mean(dim=0)}")
    print(f"- 标准差: {batch.pos.std(dim=0)}")
    
    # 打印类别分布
    unique_classes, counts = torch.unique(batch.y, return_counts=True)
    print(f"\n类别分布:")
    for cls, count in zip(unique_classes.tolist(), counts.tolist()):
        print(f"- 类别 {cls}: {count}个样本")

if __name__ == "__main__":
    print("加载ModelNet数据集...")
    train_dataset, test_dataset = load_modelnet_dataset()
    
    print("\n创建数据加载器...")
    train_loader = create_dataloader(train_dataset, batch_size=32)
    test_loader = create_dataloader(test_dataset, batch_size=32)
    
    print("\n获取一个批次并打印信息...")
    for batch in train_loader:
        print_batch_info(batch)
        break  # 只打印第一个批次
    
    print("\n可视化部分样本...")
    visualize_modelnet_samples(train_dataset, num_samples=3)
    
    print("完成!") 